# server_fixed.py
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.utils import new_agent_text_message
from a2a.server.apps import A2AStarletteApplication
from a2a.types import AgentCard, AgentSkill, AgentCapabilities
from starlette.responses import JSONResponse
from starlette.requests import Request
import json


class HelloWorldAgent:
    def say_hello(self, name: str = None):
        if name:
            return f"Hello, {name}! Welcome to A2A World!"
        else:
            return "Hello, World! Welcome to A2A World!"


class HelloWorldExecutor(AgentExecutor):
    def __init__(self):
        self.agent = HelloWorldAgent()

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        # 从上下文中获取用户输入
        user_input = context.request.input.text if context.request.input and context.request.input.text else ""

        # 简单的输入解析
        if user_input.startswith("hello ") and len(user_input) > 6:
            name = user_input[6:].strip()
            result = self.agent.say_hello(name)
        else:
            result = self.agent.say_hello()

        await event_queue.enqueue_event(new_agent_text_message(result))

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        raise Exception('not supported')


# 自定义HTTP处理器来处理请求
async def http_handler(request: Request):
    if request.method == "GET":
        # 返回agent信息
        return JSONResponse({
            "name": agent_card.name,
            "description": agent_card.description,
            "version": agent_card.version,
            "status": "running"
        })

    elif request.method == "POST":
        try:
            # 解析请求体
            body = await request.json()

            # 创建执行上下文
            from a2a.types import ExecutionRequest, Session, Input
            from a2a.server.agent_execution import RequestContext

            input_text = body.get("input", {}).get("text", "")
            session_id = body.get("session", {}).get("id", "default-session")

            # 创建请求上下文
            execution_request = ExecutionRequest(
                input=Input(text=input_text),
                session=Session(id=session_id)
            )

            context = RequestContext(request=execution_request)

            # 创建事件队列并执行
            event_queue = EventQueue()
            await executor.execute(context, event_queue)

            # 获取结果
            events = []
            while not event_queue.is_empty():
                event = await event_queue.dequeue_event()
                events.append({
                    "type": event.type,
                    "text": getattr(event, 'text', ''),
                    "timestamp": getattr(event, 'timestamp', '')
                })

            return JSONResponse({
                "events": events,
                "session_id": session_id
            })

        except Exception as e:
            return JSONResponse({"error": str(e)}, status_code=500)

    else:
        return JSONResponse({"error": "Method not allowed"}, status_code=405)


skill = AgentSkill(
    id="hello-world",
    name="Hello World",
    description="A simple hello world agent that greets users",
    tags=["greeting", "hello"],
    examples=["hello", "hello John"]
)

agent_card = AgentCard(
    name='HelloWorldAgent',
    description='A simple hello world agent for demonstration',
    version='1.0',
    defaultInputModes=['text'],
    defaultOutputModes=['text'],
    capabilities=AgentCapabilities(stream=True, tools=True),
    skills=[skill],
    url="http://0.0.0.0:9000",
)

# 构建服务器
executor = HelloWorldExecutor()
server = A2AStarletteApplication(
    agent_card=agent_card,
    http_handler=http_handler  # 使用自定义HTTP处理器
)

if __name__ == "__main__":
    import uvicorn

    print("启动A2A服务器在 http://0.0.0.0:9000")
    uvicorn.run(server.build(), host="0.0.0.0", port=9000, log_level="info")